from matplotlib import pyplot as plt
import matplotlib as mpl
from sklearn import metrics

class mdp_data():
    labels = []
    predict_prob = []
    auc = []
    macro = []
    macro_recall = []
    weighted = []
    name = ['随机森林','深度神经网络模型','SVM模型','LDA算法','决策树算法']
    color = ['green', 'gold', 'pink','blue','brown']

    def add_data(self,y_test, y_predict, auc, macro, macro_recall, weighted,auc_ave=0, g_mean_ave=0, balance_ave=0):
        self.labels.append(y_test)
        self.predict_prob.append(y_predict)
        self.auc.append(auc)
        self.macro.append(macro)
        self.macro_recall.append(macro_recall)
        self.weighted.append(weighted)

    def plot_roc_all(self):
        # 解决中文乱码和正负号问题
        mpl.rcParams["font.sans-serif"] = ["SimHei"]
        mpl.rcParams["axes.unicode_minus"] = False

        figure, axes = plt.subplots(ncols=1, nrows=2, figsize=(7.5, 6.5), dpi=100)
        ax1 = axes[0]
        ax2 = axes[1]

        plt.sca(ax1)

        for i in range(len(self.labels)):
            false_positive_rate, true_positive_rate, thresholds = metrics.roc_curve(self.labels[i], self.predict_prob[i])  # 真阳性，假阳性，阈值
            roc_auc = metrics.auc(false_positive_rate, true_positive_rate)  # 计算AUC值
            print(self.name[i]+' AUC=' + str(roc_auc))
            plt.plot(false_positive_rate, true_positive_rate, self.color[i], label=self.name[i])
        plt.legend(loc='lower right')
        plt.title('ROC')
        plt.plot([0, 1], [0, 1], 'r--')
        plt.ylabel('TPR（真阳性率）')
        plt.xlabel('FPR（伪阳性率）')

        plt.sca(ax2)
        plt.axis('off')
        plt.title('模型评价指标',loc='center', y=-0.1)
        ax2.set_position([0.16, 0.1, 0.8, 0.4])
        col_labels = ['准确率', '精确率', '召回率', 'f1值']
        row_labels = ['期望', '随机森林','深度神经网络','SVM','LDA算法','决策树算法']
        table_vals = [[0.9, 0.8, 0.75, 0.8]]
        for i in range(len(self.auc)):
            value = [self.auc[i], self.macro[i], self.macro_recall[i], self.weighted[i]]
            table_vals.append(value)
        row_colors = ['red', 'green', 'gold', 'pink','blue','brown']
        table = plt.table(cellText=table_vals, colWidths=[0.15 for x in col_labels],
                          rowLabels=row_labels, colLabels=col_labels,
                          rowColours=row_colors, colColours=row_colors,
                          loc="center")
        table.set_fontsize(14)
        table.scale(1.5, 1.5)

        plt.show()

